import numpy as np
from scipy.fft import dct, idct
import matplotlib.pyplot as plt

# 给定的数据
data = np.array([ 2, 1, -2,-7.1, 10,-2,2,9, 5, -5, -7,-2, 0, 2, 5, -5, 2, 1, -2, 2, -2, 0, 2,2, 1, -2,-7.1, 10,-2,2,9, 5, -5, -7,-2, 0, 2, 5, -5, 2, 1, -2, 2, -2, 0, 2,2, 1, -2,-7.1, 10,-2,2,9, 5, -5, -7,-2, 0, 2, 5, -5, 2, 1, -2, 2, -2, 0, 2,2, 1, -2,-7.1, 10,-2,2,9, 5, -5, -7,-2, 0, 2, 5, -5, 2, 1, -2, 2, -2, 0, 2])

print("原始数据:", data)

# 第一次 DCT 变换
dct_data = dct(data)
print("第一次 DCT 的数据:", dct_data)

watermark_dct_frame_idnex = 20
# 取整 dct_data[0]
rounded_dct_data_0 = int(dct_data[watermark_dct_frame_idnex])

# 获取最后一个比特位
last_bit = rounded_dct_data_0 & 1

# 假设水印信息是 1
watermark = 1

# 构造新的 dct_data[0]
if last_bit == 0 and watermark == 1:
    dct_data[watermark_dct_frame_idnex] = rounded_dct_data_0 + 1
elif last_bit == 1 and watermark == 0:
    dct_data[watermark_dct_frame_idnex] = rounded_dct_data_0 - 1

print("第一次 DCT 变换后加入水印的数据:", dct_data)

# IDCT 变换
idct_data = idct(dct_data)
print("IDCT 变换后的数据:", idct_data)

# 计算差值
diff = data- idct_data

# 设置图片清晰度
plt.rcParams['figure.dpi'] = 200

# 设置柱状图的位置和宽度
x = np.arange(len(data))
width = 0.35

# 创建画布和双Y坐标轴
fig, ax1 = plt.subplots()
ax2 = ax1.twinx()

# 绘制 data 和 idct_data 的柱状图
rects1 = ax1.bar(x - width / 2, data, width, label='Original Data')
rects2 = ax1.bar(x + width / 2, idct_data, width, label='IDCT Data')

# 绘制差值的折线图
line = ax2.plot(x, diff, label='Difference', color='red', marker='o')

# 设置坐标轴标签和标题
ax1.set_xlabel('Index')
ax1.set_ylabel('Value')
ax2.set_ylabel('Difference')
plt.title('Comparison of Original Data and IDCT Data with Difference')

# 添加图例
lines, labels = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines + lines2, labels + labels2, loc='upper right')

# 设置 x 轴刻度标签
ax1.set_xticks(x)

# 显示图形
plt.show()
